# coding=utf-8

from osgeo import gdal
import numpy as np
from PIL import Image
import cv2 
from matplotlib import pyplot as plt

class Dataset:
    def __init__(self, in_file):
        self.in_file = in_file  # Tiff或者ENVI文件

        dataset = gdal.Open(self.in_file)
        self.XSize = dataset.RasterXSize  # 网格的X轴像素数量
        self.YSize = dataset.RasterYSize  # 网格的Y轴像素数量
        self.GeoTransform = dataset.GetGeoTransform()  # 投影转换信息
        self.ProjectionInfo = dataset.GetProjection()  # 投影信息
        self.bands_num = dataset.RasterCount  # 波段数
        
    def get_data(self, band):
        """
        band: 读取第几个通道的数据
        """
        dataset = gdal.Open(self.in_file)
        band = dataset.GetRasterBand(band)
        data = band.ReadAsArray()
        return data

    def get_lon_lat(self):
        """
        获取经纬度信息
        """
        gtf = self.GeoTransform
        x_range = range(0, self.XSize)
        y_range = range(0, self.YSize)
        x, y = np.meshgrid(x_range, y_range)
        lon = gtf[0] + x * gtf[1] + y * gtf[2]
        lat = gtf[3] + x * gtf[4] + y * gtf[5]
        return lon, lat
    
    def writeTiff(self, im_data,im_bands, path, im_height, im_width): #,im_geotrans,im_proj,path):
        if 'int8' in im_data.dtype.name:
            datatype = gdal.GDT_Byte
        elif 'int16' in im_data.dtype.name:
            datatype = gdal.GDT_UInt16
        else:
            datatype = gdal.GDT_Float32
    
        if len(im_data.shape) == 3:
            im_bands, im_height, im_width = im_data.shape
        elif len(im_data.shape) == 2:
            im_data = np.array([im_data])
        else:
            im_bands, (im_height, im_width) = 1,im_data.shape
            
            #创建文件
        driver = gdal.GetDriverByName("GTiff")
        tifdata = driver.Create(path, im_width, im_height, im_bands, datatype)
        
        if(tifdata!= None):
            tifdata.SetGeoTransform(self.GeoTransform) #写入仿射变换参数
            tifdata.SetProjection(self.ProjectionInfo) #写入投影
            
        for i in range(im_bands):
            tifdata.GetRasterBand(i+1).WriteArray(im_data[i])
        
        #创建jpg
        if datatype != gdal.GDT_Float32:
            jpgdriver = gdal.GetDriverByName("JPEG")
            jpgdriver.CreateCopy(path[:-4]+'_gdal.jpg',tifdata)
            
        ##金字塔切片
        #gdal.SetConfigOption('HFA_USE_RlRD', 'YES')
        #tifdata.BuildOverviews(overviewlist=[2,4, 8,16,32,64,128])
        
        del tifdata
        
    def getNDVI(self, path):
        band3 = self.get_data(3) #红 660nm
        print(band3.shape)
        band5 = self.get_data(5) #近红外 840nm
        
        NDVI = (band5 - band3)/(band5 + band3)
        
        (im_height, im_width) = NDVI.shape
        self.writeTiff(NDVI, 1, path, im_height, im_width)
        
        self.buildColorNDVI2(NDVI, path[:-4] + '_color.tif')
        
    def buildColorNDVI(self, ndvi, colorfile):
        # 创建一个尺寸大小和NDVI相同的3波段图像  
        rgb = np.zeros((3, self.YSize, self.XSize), np.uint8)
        
        big = np.nanmax(ndvi)
        small = np.nanmin(ndvi)
        #NDVI着色区间
        #-0.4,-0.12,0.16,0.44,0.72
        #RG(250,0),(250,100),(250,200),(200,250),(100,250),(0,250)
        green = np.trunc((ndvi - small)/(big - small) * 7)*36 # 取整
        red = np.abs(green - 250)
        
        rgb[0] = red
        rgb[1] = green
        
        
        self.writeTiff(rgb, 3, colorfile, self.YSize, self.XSize)
        
    def buildColorNDVI2(self, ndvi, colorfile):
        # 创建一个尺寸大小和NDVI相同的3波段图像  
        rgb = np.zeros((3, self.YSize, self.XSize), np.uint8)
        
        #NDVI着色区间
        #-0.4,-0.12,0.16,0.44,0.72
        #RG(250,0),(250,100),(250,200),(200,250),(100,250),(0,250)
        #很差，差，一般，较好，良好，很好
        
        
        red = ndvi.copy()
        red[red <= 0.16] = 250
        red[red <= 0.44] = 200
        red[red <= 0.72] = 100
        red[red < 100] = 0
        
        green = ndvi.copy()
        green[green >= 0.16] = 250
        green[green <= -0.4] = 0
        green[green <= -0.12] = 100
        green[green <= 0.16] = 200
        
        rgb[0] = red
        rgb[1] = green
        
        
        self.writeTiff(rgb, 3, colorfile, self.YSize, self.XSize)
        
        #计算统计
        calcTable = {}
        calcTable['很差'] = np.sum(ndvi <=-0.4)
        calcTable['差'] = np.sum(np.logical_and(ndvi > -0.4, ndvi <= -0.12))
        calcTable['一般'] = np.sum(np.logical_and(ndvi > -0.12, ndvi <= 0.16))
        calcTable['较好'] = np.sum(np.logical_and(ndvi > 0.16, ndvi <= 0.44))
        calcTable['良好'] = np.sum(np.logical_and(ndvi > 0.44, ndvi <= 0.72))
        calcTable['很好'] = np.sum(ndvi > 0.72)
        
        print(calcTable)
        ##get jpg
        #im = Image.open(colorfile)
        #im.save(colorfile[:-4] + '.jpg')
        
    def calcArea(self, f):
        img = cv2.imread(f, 0)
        color = ('b','g','r')
        
        #for i,col in enumerate(color):
        histr = cv2.calcHist([img],[0],None,[256],[0,256])
        plt.plot(histr)
        plt.xlim([0,256])
        plt.show()
        

# 以下代码演示读取E:/data/dataset.tif的第一个通道的数据，并且获取经纬度信息
filename = 'Z:/CloudStation/GIS/项目/安州智慧农业/多光谱/agri/pj-RadCalibreation.tif'
dataset = Dataset(filename)

dataset.getNDVI('./ndvi/ndvi_test2.tif')
longitude, latitude = dataset.get_lon_lat()  # 获取经纬度信息
print(dataset.XSize)
print(dataset.YSize)
print(dataset.bands_num)
dataset.calcArea('./ndvi/ndvi_test1_color.tif')
